//===----------------------------------------------------------------------===//
//
//                         BusTub
//
// hash_join_executor.cpp
//
// Identification: src/execution/hash_join_executor.cpp
//
// Copyright (c) 2015-2021, Carnegie Mellon University Database Group
//
//===----------------------------------------------------------------------===//

#include "execution/executors/hash_join_executor.h"
#include "type/value_factory.h"

namespace bustub {

HashJoinExecutor::HashJoinExecutor(ExecutorContext *exec_ctx, const HashJoinPlanNode *plan,
                                   std::unique_ptr<AbstractExecutor> &&left_child,
                                   std::unique_ptr<AbstractExecutor> &&right_child)
    : AbstractExecutor(exec_ctx), plan_(plan), lch_(std::move(left_child)), rch_(std::move(right_child)) {
  if (!(plan->GetJoinType() == JoinType::LEFT || plan->GetJoinType() == JoinType::INNER)) {
    // Note for 2023 Spring: You ONLY need to implement left join and inner join.
    throw bustub::NotImplementedException(fmt::format("join type {} not supported", plan->GetJoinType()));
  }
}

void HashJoinExecutor::Init() {
  // std::cout << plan_->ToString(true) << std::endl;
  ht_.Clear();
  lch_->Init();
  rch_->Init();
  Tuple tuple;
  RID rid;
  const Schema &rschema = rch_->GetOutputSchema();

  const JoinType &join_type = plan_->GetJoinType();
  if (join_type == JoinType::LEFT || join_type == JoinType::INNER) {
    while (rch_->Next(&tuple, &rid)) {
      ht_.Insert(tuple, rschema, plan_->RightJoinKeyExpressions());
    }
  }
  cur_idx_ = 0;
  left_tuple_ = Tuple{};
  cur_join_value_ = JoinValue{};
}
auto HashJoinExecutor::Next(Tuple *tuple, RID *rid) -> bool {
  const JoinType &join_type = plan_->GetJoinType();
  if (join_type == JoinType::LEFT) {
    return JoinLeft(tuple);
  } else if (join_type == JoinType::INNER) {
    return JoinInner(tuple);
  }
  return false;
}

auto HashJoinExecutor::JoinLeft(Tuple *tuple) -> bool {
  const Schema &lschema = lch_->GetOutputSchema();
  RID lrid;
  while (true) {
    if (cur_idx_ == cur_join_value_.vals.size()) {
      if (not has_match_ && left_tuple_.GetLength() != 0) {
        has_match_ = true;
        GenerateNullTuple(tuple);
        return true;
      }
      if (not lch_->Next(&left_tuple_, &lrid)) {
        return false;
      }
      if (not ht_.GetValue(left_tuple_, lschema, plan_->LeftJoinKeyExpressions(), cur_join_value_)) {
        has_match_ = true;
        GenerateNullTuple(tuple);
        return true;
      }
      has_match_ = false;
      cur_idx_ = 0;
    }
    const auto &[hash_key, r_tuple] = cur_join_value_.vals[cur_idx_];
    ++cur_idx_;
    if (hash_key == JoinHashTable::GetHashKey(left_tuple_, lschema, plan_->LeftJoinKeyExpressions())) {
      has_match_ = true;
      GenerateGoodTuple(tuple, r_tuple);
      return true;
    }
  }
}

auto HashJoinExecutor::JoinInner(Tuple *tuple) -> bool {
  const Schema &lschema = lch_->GetOutputSchema();
  RID lrid;
  while (true) {
    if (cur_idx_ == cur_join_value_.vals.size()) {
      if (not lch_->Next(&left_tuple_, &lrid)) {
        break;
      }
      if (not ht_.GetValue(left_tuple_, lschema, plan_->LeftJoinKeyExpressions(), cur_join_value_)) {
        continue;
      }
      cur_idx_ = 0;
    }
    const auto &[hash_key, r_tuple] = cur_join_value_.vals[cur_idx_];
    ++cur_idx_;
    if (hash_key == JoinHashTable::GetHashKey(left_tuple_, lschema, plan_->LeftJoinKeyExpressions())) {
      GenerateGoodTuple(tuple, r_tuple);
      return true;
    }
  }
  return false;
}

void HashJoinExecutor::GenerateNullTuple(Tuple *tuple) {
  const Schema &lschema = lch_->GetOutputSchema();
  const Schema &rschema = rch_->GetOutputSchema();
  std::vector<Value> vals;
  vals.reserve(lschema.GetColumnCount() + rschema.GetColumnCount());
  for (uint32_t i = 0; i < lschema.GetColumnCount(); ++i) {
    vals.emplace_back(left_tuple_.GetValue(&lschema, i));
  }
  for (uint32_t i = 0; i < rschema.GetColumnCount(); ++i) {
    vals.emplace_back(ValueFactory::GetNullValueByType(rschema.GetColumn(i).GetType()));
  }
  Tuple ret(vals, &GetOutputSchema());
  *tuple = std::move(ret);
}

void HashJoinExecutor::GenerateGoodTuple(Tuple *tuple, const Tuple &right_tuple) {
  const Schema &lschema = lch_->GetOutputSchema();
  const Schema &rschema = rch_->GetOutputSchema();
  std::vector<Value> vals;
  vals.reserve(lschema.GetColumnCount() + rschema.GetColumnCount());
  for (uint32_t i = 0; i < lschema.GetColumnCount(); ++i) {
    vals.emplace_back(left_tuple_.GetValue(&lschema, i));
  }
  for (uint32_t i = 0; i < rschema.GetColumnCount(); ++i) {
    vals.emplace_back(right_tuple.GetValue(&rschema, i));
  }
  Tuple ret(vals, &GetOutputSchema());
  *tuple = std::move(ret);
}
}  // namespace bustub
